package util

import (
	"chat/config"
	"context"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

var Mongo *mongoPool

type mongoPool struct {
	pool    chan *mongo.Client
	maxSize int
	url     string
	ctx     context.Context
}

func MongoInit() {
	Mongo = newMongoPool(config.Config.Mongo)
}

func newMongoPool(mongoConfig config.MongoConfig) *mongoPool {
	return &mongoPool{
		pool:    make(chan *mongo.Client, mongoConfig.MaxSize),
		maxSize: mongoConfig.MaxSize,
		url:     mongoConfig.Url,
		ctx:     context.Background(),
	}
}

func (p *mongoPool) Get() (*mongo.Client, error) {
	select {
	case conn := <-p.pool:
		return conn, nil
	default:
		client, err := mongo.Connect(p.ctx, options.Client().ApplyURI(p.url))
		if err != nil {
			return nil, err
		}
		if err != nil {
			return nil, err
		}
		return client, nil
	}
}

func (p *mongoPool) Put(conn *mongo.Client) {
	select {
	case p.pool <- conn:
	default:
		conn.Disconnect(p.ctx)
	}
}

func (p *mongoPool) GetChatCollection(collection string) (*mongo.Collection, error) {
	c, err := p.Get()
	if err != nil {
		return nil, err
	}
	return c.Database("chat").Collection(collection), nil
}
